package org.springmvc.mybatis.pagination;

import org.apache.ibatis.executor.parameter.ParameterHandler;
import org.apache.ibatis.executor.statement.RoutingStatementHandler;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ParameterMapping;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.scripting.defaults.DefaultParameterHandler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.List;
import java.util.Map;
import java.util.Properties;

/**
 * MyBatis分页拦截器
 *
 * @author tangjialin on 2018-08-17.
 */
@Intercepts({@Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})})
public class StatementHandlerInterceptor implements Interceptor {

    private Logger logger = LoggerFactory.getLogger(getClass());

    private PageableHandler handler;

    public StatementHandlerInterceptor() {
    }

    public StatementHandlerInterceptor(PageableHandler handler) {
        this.handler = handler;
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        RoutingStatementHandler handler = (RoutingStatementHandler) invocation.getTarget();
        BoundSql boundSql = handler.getBoundSql();
        Object obj = boundSql.getParameterObject();
        Pageable page = null;
        if (obj instanceof Pageable) {
            // 这里我们简单的通过传入的是Page对象就认定它是需要进行分页操作的。
            page = (Pageable) obj;
        } else if (obj instanceof Map) {
            // 多参数,包含Page对象
            /** 此处的obj是一个Map,具体的实现类看这里:{@link org.apache.ibatis.binding.MapperMethod.ParamMap}
             * 但是这个类的get方法实现不太科学,导致此处判断不能直接使用get方法,否则会抛出异常 */
            Map<String, Object> paramMap = (Map<String, Object>) obj;
            page = paramMap.containsKey("pageable") ? (Pageable) paramMap.get("pageable") : null;
        }
        if (page != null) {
            StatementHandler delegate = (StatementHandler) ReflectUtils.getFieldValue(handler, "delegate");
            MappedStatement mappedStatement = (MappedStatement) ReflectUtils.getFieldValue(delegate, "mappedStatement");
            Connection connection = (Connection) invocation.getArgs()[0];
            page.setTotal(this.getRecordTotal(obj, mappedStatement, connection));
            String pageSql = this.getPageSql(page, boundSql.getSql());
            ReflectUtils.setFieldValue(boundSql, "sql", pageSql);
        }
        // 将执行权交给下一个拦截器
        return invocation.proceed();
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {
        if (logger.isInfoEnabled()) { logger.info("setProperties"); }
    }

    /**
     * 获取数据库的分页查询语句
     *
     * @param pageable 分页对象
     * @param sql      原SQL语句
     * @return 返回数据库分页查询语句
     */
    private String getPageSql(Pageable pageable, String sql) {
        return handler.handler(pageable, sql);
    }

    /**
     * 获取记录总数
     *
     * @param data            Mapper映射语句对应的原始参数对象
     * @param mappedStatement Mapper映射语句
     * @param connection      当前的数据库连接
     */
    private long getRecordTotal(Object data, MappedStatement mappedStatement, Connection connection) throws SQLException {
        BoundSql boundSql = mappedStatement.getBoundSql(data);
        String countSql = this.getCountSql(boundSql.getSql());
        List<ParameterMapping> parameterMappings = boundSql.getParameterMappings();
        BoundSql countBoundSql = new BoundSql(mappedStatement.getConfiguration(), countSql, parameterMappings, data);
        ParameterHandler parameterHandler = new DefaultParameterHandler(mappedStatement, data, countBoundSql);
        PreparedStatement pstmt = null;
        ResultSet rs = null;
        try {
            pstmt = connection.prepareStatement(countSql);
            parameterHandler.setParameters(pstmt);
            // 执行记录总条数查询
            rs = pstmt.executeQuery();
            return rs.next() ? rs.getLong(1) : 0L;
        } finally {
            try { if (rs != null) { rs.close(); } } catch (SQLException e) {}
            try { if (pstmt != null) { pstmt.close(); } } catch (SQLException e) {}
        }
    }

    /**
     * 根据原SQL语句获取对应的查询总记录数的SQL语句
     *
     * @param sql
     * @return
     */
    private String getCountSql(String sql) {
        return "SELECT COUNT(1) FROM (" + sql + ") temp";
    }

    public PageableHandler getHandler() {
        return handler;
    }

    public StatementHandlerInterceptor setHandler(PageableHandler handler) {
        this.handler = handler;
        return this;
    }

    /**
     * 分页处理
     *
     * @author tangjialin on 2018-08-22.
     */
    public interface PageableHandler {
        /**
         * 获取分页SQL语句
         *
         * @param pageable 分页参数
         * @param sql      查询语句
         * @return 返回分页查询语句
         */
        String handler(Pageable pageable, String sql);
    }

    /**
     * ORACLE数据库分页处理
     *
     * @author tangjialin on 2018-08-22.
     */
    public static class OraclePageableHandler implements PageableHandler {
        @Override
        public String handler(Pageable pageable, String sql) {
            int offset = (pageable.getPageNumber() - 1) * pageable.getPageSize() + 1;
            return new StringBuilder(sql.length() + 256)
                    .append("SELECT * FROM (")
                    .append("SELECT TB.*, ROWNUM RN FROM (")
                    .append(sql)
                    .append(") TB WHERE ROWNUM < ").append(offset + pageable.getPageSize())
                    .append(") WHERE RN >= ").append(offset)
                    .toString();
        }
    }

    /**
     * MySql数据库分页处理
     *
     * @author tangjialin on 2018-08-22.
     */
    public static class MySqlPageableHandler implements PageableHandler {
        @Override
        public String handler(Pageable pageable, String sql) {
            int offset = (pageable.getPageNumber() - 1) * pageable.getPageSize();
            return new StringBuilder(sql.length() + 64)
                    .append(sql).append(" LIMIT ").append(offset).append(",").append(pageable.getPageSize())
                    .toString();
        }
    }

    /**
     * MSSql数据库分页处理
     *
     * @author tangjialin on 2018-08-22.
     */
    public static class MsSqlPageableHandler implements PageableHandler {
        @Override
        public String handler(Pageable pageable, String sql) {
            int pageNumber = pageable.getPageNumber();
            int offset = pageNumber * pageable.getPageSize();
            int insert = 6;
            StringBuilder sqlBuilder = new StringBuilder(sql.length() + 256)
                    .append(sql.trim()).insert(insert, ") ").insert(insert, offset).insert(insert, " TOP(");
            if (pageNumber > 1) {
                sqlBuilder.insert(0, "SELECT query.* FROM (SELECT inner_query.*, ROW_NUMBER() OVER (ORDER BY CURRENT_TIMESTAMP) as _rn_ FROM ( ")
                        .append(" ) as inner_query ) as query WHERE _rn_ > ").append(offset - pageable.getPageSize())
                ;
            }
            return sqlBuilder.toString();
        }
    }
}